from sklearn.preprocessing import MinMaxScaler
import akshare as ak
import os
import pandas as pd
import math
import torch
from torch.utils.data import Dataset, DataLoader, TensorDataset
import torch.nn as nn
from tqdm import trange

start_data='20100101'
end_data='20240507'
code='000001'
data_file=code+'.csv'
if os.path.exists(data_file):
    df=pd.read_csv(data_file)
else:
    df = ak.stock_zh_a_hist(symbol=code, period="daily", start_date=start_data, end_date=end_data, adjust="")
    df.to_csv(data_file)
df.rename(columns={'日期': 'Date', '开盘': 'Open', '最高': 'High', '最低': 'Low', '收盘': 'Close', '成交量': 'Vol'}, inplace=True)
df=df[["Date", "Open", "High", "Low", "Close", "Vol"]]
print(df.head)

train_size=math.floor(df.shape[0]*0.8)
test_size=df.shape[0]-train_size

mms = MinMaxScaler()
# 后800条记录用于提取测试集
mms.fit(df.iloc[:-800][["Open", "High", "Low", "Close", "Vol"]])
df[["Open", "High", "Low", "Close", "Vol"]] = mms.transform(df[["Open", "High", "Low", "Close", "Vol"]])
print(df.shape)

X = []
y = []
# 时间长度设定为30
for i in range(30, len(df)):
    X.append(df.iloc[i-30:i, 1:6].values)  # 输入数据未取到i时刻
    y.append(df.iloc[i, 4])  # 标签数据为i时刻
X = torch.tensor(X, dtype=torch.float)
y = torch.tensor(y, dtype=torch.float).view(-1, 1)
X.shape, y.shape
## 输出
(torch.Size([3997, 30, 5]), torch.Size([3997, 1]))

N = -800
X_train, X_test = X[:N], X[N:]
y_train, y_test = y[:N], y[N:]
trainloader = DataLoader(TensorDataset(X_train, y_train), 64, True)
X_train.shape, X_test.shape
## 输出
(torch.Size([3197, 30, 5]), torch.Size([800, 30, 5]))

class ModelDNN(nn.Module):
    def __init__(self, input_size, hiddens=[64, 32, 8]):
        super().__init__()
        self.hiddens = hiddens
        self.net = nn.Sequential(nn.Flatten())
        for pre, nxt in zip([input_size]+hiddens[:-1], hiddens):
            self.net.append(nn.Linear(pre, nxt))
            self.net.append(nn.ReLU())
        self.net.append(nn.Linear(hiddens[-1], 1))
        
    def forward(self, x):
        return self.net(x)

class ModelCNN(nn.Module):
    def __init__(self, in_channels, hidden_channels=[8, 4]):
        # input: N x C x L
        # C: 5->8->4
        # L: 30->28->14->12->6
        super().__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.net = nn.Sequential()
        for in_channels, out_channels in zip([in_channels]+hidden_channels[:-1], hidden_channels):
            self.net.append(nn.Conv1d(in_channels, out_channels, kernel_size=3))
            self.net.append(nn.AvgPool1d(kernel_size=2))
        self.net.append(nn.Flatten())
        self.net.append(nn.Linear(6*hidden_channels[-1], 1))
        
    def forward(self, x):
        x = x.permute(0, 2, 1)
        return self.net(x)

class ModelRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # 注意，这里设置batch_first=True
        self.gru = nn.GRU(input_size=input_size, hidden_size=hidden_size, batch_first=True)
        self.activation = nn.ReLU()
        self.output = nn.Linear(hidden_size, 1)
        
    def forward(self, x):
        _, hidden = self.gru(x)
        hidden = hidden.squeeze(0)
        hidden = self.activation(hidden)
        return self.output(hidden)


modelDNN = ModelDNN(30*5)
optimizer = torch.optim.Adam(modelDNN.parameters())
criterion = nn.MSELoss()
for i in trange(100):
    for X, y in trainloader:
        y_pred = modelDNN(X)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if (i+1) % 10 ==0:
        with torch.no_grad():
            train_pred = modelDNN(X_train)
            train_mse = criterion(train_pred, y_train)
            test_pred = modelDNN(X_test)
            test_mse = criterion(test_pred, y_test)
            print(i, train_mse.item(), test_mse.item())


modelCNN = ModelCNN(5)
optimizer = torch.optim.Adam(modelCNN.parameters())
criterion = nn.MSELoss()
for i in trange(100):
    for X, y in trainloader:
        y_pred = modelCNN(X)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if (i+1) % 10 ==0:
        with torch.no_grad():
            train_pred = modelCNN(X_train)
            train_mse = criterion(train_pred, y_train)
            test_pred = modelCNN(X_test)
            test_mse = criterion(test_pred, y_test)
            print(i, train_mse.item(), test_mse.item())


modelRNN = ModelRNN(5, 4)
optimizer = torch.optim.Adam(modelRNN.parameters())
criterion = nn.MSELoss()
for i in trange(100):
    for X, y in trainloader:
        y_pred = modelRNN(X)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if (i+1) % 10 ==0:
        with torch.no_grad():
            train_pred = modelRNN(X_train)
            train_mse = criterion(train_pred, y_train)
            test_pred = modelRNN(X_test)
            test_mse = criterion(test_pred, y_test)
            print(i, train_mse.item(), test_mse.item())
